#include "MrtProcess.h"
#include "Utils/LogUtil.h"
#include "Render/RenderState.h"
#include "Core/CoreGlobals.h"
#include "Game/Scene/Entity.h"
#include "Game/Component/Model.h"
#include "Render/Buffer/VertexBuffer.h"
#include "Render/Buffer/IndexBuffer.h"
#include "Render/Buffer/CustomBuffer.h"
#include "LocalShader.h"

void MrtProcess::OnInit()
{
	mModelEntitys = mGameScene->GetEntitysByTypes<LocalEntity, Model, TextureMaterial>();

	std::vector<Ubo> ubos;
	ubos.resize(mModelEntitys.size());
	for (uint32 i = 0; i < mModelEntitys.size(); ++i)
	{
		ubos[i].proj = mGameScene->GetCamera().GetProj();
		ubos[i].view = mGameScene->GetCamera().GetView();
		ubos[i].model = mModelEntitys[i]->GetTransform().GetLocalMatrix();
	}

	float cubeSize = mGameScene->GetCamera().GetFarZ();

	UboSkybox uboSkybox;
	uboSkybox.proj = mGameScene->GetCamera().GetProj();
	uboSkybox.model = mGameScene->GetCamera().GetView();
	// Dont not calc translate, the far always not change
	uboSkybox.model[3] = glm::vec4(0.0f, 0.0f, 0.0f, 1.0f);

	ENQUEUE_RENDER_COMMAND(BuildModel)([this, uboSkybox, ubos, cubeSize](VkCommandBuffer commandBuffer)
		{

			GeometryGenerator geometryGenerator;
			GeometryGenerator::MeshData cubeData = geometryGenerator.CreateCube(cubeSize, cubeSize, cubeSize, 1);

			std::vector<glm::vec3> vertexs;
			vertexs.resize(cubeData.Vertices.size());
			for (uint32 i = 0; i < cubeData.Vertices.size(); ++i)
			{
				vertexs[i] = cubeData.Vertices[i].Position;
			}

			mSkyboxVertexBuffer = VertexBuffer::Builder()
				.SetData(vertexs.data())
				.SetSize(vertexs.size() * sizeof(glm::vec3))
				.Build();

			mSkyboxIndexBuffer = IndexBuffer::Builder()
				.SetData(cubeData.Indices32.data())
				.SetSize(cubeData.Indices32.size() * sizeof(uint32))
				.SetIndexCount(cubeData.Indices32.size())
				.Build();

			UboSkybox uboSkyboxData = uboSkybox;

			mSkyboxUboBuffer = CustomBuffer::Builder()
				.SetData(&uboSkyboxData)
				.SetSize(sizeof(UboSkybox))
				.SetMemoryProperty(VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
				.Build();


			mSkyboxTexture = CubemapTexture::Builder()
				.SetAutoMip()
				.SetPaths({
				"Hdr\\TreeSky_nz.hdr",
				"Hdr\\TreeSky_pz.hdr",
				"Hdr\\TreeSky_py.hdr",
				"Hdr\\TreeSky_ny.hdr",
				"Hdr\\TreeSky_px.hdr",
				"Hdr\\TreeSky_nx.hdr"
					})
				.SetFormat(VK_FORMAT_R32G32B32A32_SFLOAT)
				.Build();

			mSkyboxSampler = RHISampler::Builder()
				.SetAddressModeU(VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE)
				.SetAddressModeV(VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE)
				.SetAddressModeW(VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE)
				.SetCompareOp(VK_COMPARE_OP_NEVER)
				.SetMaxLod(mSkyboxTexture->GetMipLevels())
				.SetBorderColor(VK_BORDER_COLOR_FLOAT_OPAQUE_WHITE)
				.Build();

			mSkyboxDescriptor = RHIDescriptor::Builder()
				.PushShader<SkyboxVertexShader>()
				.PushShader<SkyboxFragmentShader>()
				.SetDescriptorCount(0, 1)
				.Build();

			mSkyboxDescriptor->SetValue("ubo", mSkyboxUboBuffer);
			mSkyboxDescriptor->SetValue("samplerCubeMap", mSkyboxTexture, mSkyboxSampler);
			mSkyboxDescriptorMask = mSkyboxDescriptor->AllocateDescriptor();









			std::vector<Ubo> entityUbos = ubos;

			for (uint32 i = 0; i < mModelEntitys.size(); ++i)
			{
				Model& model = mModelEntitys[i]->GetComponent<Model>();

				mModelEntitys[i]->mVertexBuffer = VertexBuffer::Builder()
					.SetData(model.mVertexs.data())
					.SetSize(model.mVertexs.size() * sizeof(Vertex))
					.Build();

				mModelEntitys[i]->mIndexBuffer = IndexBuffer::Builder()
					.SetData(model.mIndices.data())
					.SetSize(model.mIndices.size() * sizeof(uint32))
					.SetIndexCount(model.mIndices.size())
					.Build();

				TextureMaterial& textureMaterial = mModelEntitys[i]->GetComponent<TextureMaterial>();

				mModelEntitys[i]->mColorTexture = Picture2DTexture::Builder()
					.SetPath(textureMaterial.colorPath)
					.SetAutoMip()
					.Build();

				mModelEntitys[i]->mNormalTexture = Picture2DTexture::Builder()
					.SetPath(textureMaterial.normalPath)
					.SetAutoMip()
					.Build();

				mModelEntitys[i]->mColorSampler = RHISampler::Builder()
					.SetMaxLod(mModelEntitys[i]->mColorTexture->GetMipLevels())
					.Build();

				mModelEntitys[i]->mNormalSampler = RHISampler::Builder()
					.SetMaxLod(mModelEntitys[i]->mNormalTexture->GetMipLevels())
					.Build();

				mModelEntitys[i]->mUboBuffer = CustomBuffer::Builder()
					.SetData(&entityUbos[i])
					.SetSize(sizeof(Ubo))
					.SetMemoryProperty(VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
					.Build();
			}

			mDescriptor = RHIDescriptor::Builder()
				.PushShader<MrtVertexShader>()
				.PushShader<MrtFragmentShader>()
				.SetDescriptorCount(0, mModelEntitys.size())
				.Build();

			for (auto& entity : mModelEntitys)
			{
				mDescriptor->SetValue("ubo", entity->mUboBuffer);
				mDescriptor->SetValue("samplerColor", entity->mColorTexture, entity->mColorSampler);
				mDescriptor->SetValue("samplerNormal", entity->mNormalTexture, entity->mNormalSampler);
				entity->mDescriptorMask = mDescriptor->AllocateDescriptor();
			}

			PositionRenderTarget = ColorRenderTarget::Builder()
				.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.SetMultiple(true)
				.Build(this, "PositionRenderTarget");

			NormalRenderTarget = ColorRenderTarget::Builder()
				.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.SetMultiple(true)
				.Build(this, "NormalRenderTarget");

			AlbedoRenderTarget = ColorRenderTarget::Builder()
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.SetMultiple(true)
				.Build(this, "AlbedoRenderTarget");

			mDepthRenderTarget = DepthRenderTarget::Builder()
				.SetUsage(VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetAspectMask(VK_IMAGE_ASPECT_DEPTH_BIT)
				.SetFormat(VK_FORMAT_D32_SFLOAT_S8_UINT)
				.SetMultiple(true)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.Build(this, "DepthRenderTarget");

			mPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mDescriptor)
				.AddColorRenderTarget(PositionRenderTarget)
				.AddColorRenderTarget(NormalRenderTarget)
				.AddColorRenderTarget(AlbedoRenderTarget)
				.SetDepthRenderTarget(mDepthRenderTarget)
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.Build();

			mRenderPass = RHIRenderPass::Builder()
				.AddRenderTarget(PositionRenderTarget)
				.AddRenderTarget(NormalRenderTarget)
				.AddRenderTarget(AlbedoRenderTarget)
				.AddRenderTarget(mDepthRenderTarget)
				.AddPipeline(mPipeline)
				.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.Build();

			mSkyboxPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mSkyboxDescriptor)
				.AddColorRenderTarget(PositionRenderTarget)
				.AddColorRenderTarget(NormalRenderTarget)
				.AddColorRenderTarget(AlbedoRenderTarget)
				.SetDepthRenderTarget(mDepthRenderTarget)
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetRenderPass(mRenderPass->GetRenderPass())
				.SetCullMode(VK_CULL_MODE_BACK_BIT)
				.Build();

			mFrameBuffer = RHIFrameBuffer::Builder()
				.SetRenderPass(mRenderPass)
				.Build();

			mCommandBuffer = RHICommandBuffer::Builder()
				.SetQueue(VK_QUEUE_GRAPHICS_BIT)
				.Build(this, "OfflineCommandBuffer");

			mWaitGraphicsSemaphore = RHISemaphore::Builder()
				.SetWaitDstStage(VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT)
				.SetSignaled(true)
				.Build(this, "WaitGraphicsSemaphore");

			mSignalOfflineSemaphore = RHISemaphore::Builder()
				.SetWaitDstStage(VK_PIPELINE_STAGE_VERTEX_INPUT_BIT)
				.Build(this, "SignalOfflineSemaphore");
		});
}

void MrtProcess::OnTick(float deltaTime)
{
	std::vector<Ubo> ubos;
	ubos.resize(mModelEntitys.size());
	for (uint32 i = 0; i < mModelEntitys.size(); ++i)
	{
		ubos[i].proj = mGameScene->GetCamera().GetProj();
		ubos[i].view = mGameScene->GetCamera().GetView();
		ubos[i].model = mModelEntitys[i]->GetTransform().GetLocalMatrix();
	}

	UboSkybox uboSkybox;
	uboSkybox.proj = mGameScene->GetCamera().GetProj();
	uboSkybox.model = mGameScene->GetCamera().GetView();
	// Dont not calc translate, the far always not change
	uboSkybox.model[3] = glm::vec4(0.0f, 0.0f, 0.0f, 1.0f);

	ENQUEUE_RENDER_COMMAND(RecordCommand)([this, ubos, uboSkybox](VkCommandBuffer commandBuffer)
		{

			UboSkybox uboSkyboxData = uboSkybox;
			mSkyboxUboBuffer->WriteMemory(&uboSkyboxData);

			std::vector<Ubo> entityUbos = ubos;

			for (uint32 i = 0; i < mModelEntitys.size(); ++i)
			{
				mModelEntitys[i]->mUboBuffer->WriteMemory(&entityUbos[i]);
			}

			GRenderState->AddWaitSemaphore(mSignalOfflineSemaphore);
			GRenderState->AddSignalSemaphore(mWaitGraphicsSemaphore);

			VkCommandBuffer offlineCommandBuffer = mCommandBuffer->BeginCommandBuffer();

			VkRenderPassBeginInfo renderPassBeginInfo = mFrameBuffer->GetRenderPassBeginInfo();

			vkCmdBeginRenderPass(offlineCommandBuffer, &renderPassBeginInfo, VK_SUBPASS_CONTENTS_INLINE);

			vkCmdSetViewport(offlineCommandBuffer, 0, 1, &GRenderState->GetViewport());
			vkCmdSetScissor(offlineCommandBuffer, 0, 1, &GRenderState->GetScissor());





			// Draw Skybox
			vkCmdBindPipeline(offlineCommandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mSkyboxPipeline->GetPipeline());

			vkCmdBindDescriptorSets(offlineCommandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mSkyboxPipeline->GetPipelineLayout(), 0, 1, &mSkyboxDescriptorMask.GetDescriptorSet(), 0, nullptr);

			VkBuffer skyboxVertexBuffer = mSkyboxVertexBuffer->GetBuffer();
			VkDeviceSize skyboxOffsets = 0;
			vkCmdBindVertexBuffers(offlineCommandBuffer, 0, 1, &skyboxVertexBuffer, &skyboxOffsets);

			VkBuffer skyboxIndexBuffer = mSkyboxIndexBuffer->GetBuffer();
			vkCmdBindIndexBuffer(offlineCommandBuffer, skyboxIndexBuffer, 0, VK_INDEX_TYPE_UINT32);

			vkCmdDrawIndexed(offlineCommandBuffer, mSkyboxIndexBuffer->GetIndexCount(), 1, 0, 0, 0);






			// Draw Scene
			vkCmdBindPipeline(offlineCommandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipeline());

			for (const auto& entity : mModelEntitys)
			{
				vkCmdBindDescriptorSets(offlineCommandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipelineLayout(), 0, 1, &entity->mDescriptorMask.GetDescriptorSet(), 0, nullptr);

				VkBuffer vertexBuffer = entity->mVertexBuffer->GetBuffer();
				VkDeviceSize offsets = 0;
				vkCmdBindVertexBuffers(offlineCommandBuffer, 0, 1, &vertexBuffer, &offsets);

				VkBuffer indexBuffer = entity->mIndexBuffer->GetBuffer();
				vkCmdBindIndexBuffer(offlineCommandBuffer, indexBuffer, 0, VK_INDEX_TYPE_UINT32);

				vkCmdDrawIndexed(offlineCommandBuffer, entity->mIndexBuffer->GetIndexCount(), 1, 0, 0, 0);
			}

			vkCmdEndRenderPass(offlineCommandBuffer);

			// mCommandBuffer->FlushCommandBuffer(offlineCommandBuffer, { mWaitGraphicsSemaphore }, { mSignalOfflineSemaphore });
		});
}

void MrtProcess::OnResize()
{
	PositionRenderTarget = ColorRenderTarget::Builder()
		.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.SetMultiple(true)
		.Build(this, "PositionRenderTarget");

	NormalRenderTarget = ColorRenderTarget::Builder()
		.SetFormat(VK_FORMAT_R16G16B16A16_SFLOAT)
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.SetMultiple(true)
		.Build(this, "NormalRenderTarget");

	AlbedoRenderTarget = ColorRenderTarget::Builder()
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.SetMultiple(true)
		.Build(this, "AlbedoRenderTarget");

	mDepthRenderTarget = DepthRenderTarget::Builder()
		.SetUsage(VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetAspectMask(VK_IMAGE_ASPECT_DEPTH_BIT)
		.SetFormat(VK_FORMAT_D32_SFLOAT_S8_UINT)
		.SetMultiple(true)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.Build(this, "DepthRenderTarget");

	mFrameBuffer = RHIFrameBuffer::Builder()
		.AddRenderTarget(PositionRenderTarget)
		.AddRenderTarget(NormalRenderTarget)
		.AddRenderTarget(AlbedoRenderTarget)
		.AddRenderTarget(mDepthRenderTarget)
		.SetRenderPass(mRenderPass)
		.Build();
}
